import dspy

from dspy.datasets import MATH


qwen7b = dspy.LM(
    'ollama_chat/qwen2.5', 
    api_base='http://localhost:11434', 
    api_key='', 
    temperature=1.0,
    cache=False,
    num_retries=1000)

o1_mini = dspy.LM(
    'openai/o1-mini', 
    api_key='sk-xxxxxx',
    temperature=1.0,
    max_tokens=5000,
    cache=False,
    num_retries=1000)

dspy.configure(lm=qwen7b)

dataset = MATH(subset='algebra')
print(len(dataset.train), len(dataset.dev))

module = dspy.ChainOfThought("question -> answer")

THREADS = 24
kwargs = dict(num_threads=THREADS, display_progress=True, display_table=5)
evaluate = dspy.Evaluate(devset=dataset.dev, metric=dataset.metric, **kwargs)

kwargs = dict(num_threads=THREADS, teacher_settings=dict(lm=o1_mini), prompt_model=qwen7b)
optimizer = dspy.MIPROv2(metric=dataset.metric, auto="medium", **kwargs)

kwargs = dict(requires_permission_to_run=False, max_bootstrapped_demos=4, max_labeled_demos=4)
optimized_module = optimizer.compile(module, trainset=dataset.train, **kwargs)

optimized_module.save("optimized_module.json")

